/*
* nv_detector.h
* Created on: 20230620
* Author: sumang
* Description: armor detector of tensorrt
*/
#ifndef NV_DETECT_H_
#define NV_DETECT_H_

#ifndef USE_NVIDIA

#include <vector>
#include <opencv2/opencv.hpp>
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "logger.h"
#include "common.h"
#include "buffers.h"
#include "detector.h"
#include "common_structs.h"


#define OUTPUT_SIZE_8 1632000
#define OUTPUT_SIZE_16 408000
#define OUTPUT_SIZE_32 102000
#define OUTPUT_SIZE_G640 2142000

class Nv_Detector:public MyDetector 
{
public:
    std::vector<unsigned char> engine_data_;
    std::unique_ptr<nvinfer1::IRuntime> runtime_;
    std::shared_ptr<nvinfer1::ICudaEngine> engine_;
    std::unique_ptr<nvinfer1::IExecutionContext> context_;
    int inputIndex;
    int output_stride8;
    int output_stride16;
    int output_stride32;

    void* buffers_[5];
    cudaStream_t stream_;

    const int batch_size_ = 2;
    cv::Mat pr_img;
    cv::Mat show_left;
    cv::Mat show_right;

    // void do_merge_nms();
    void preprocess() override;
    void inference() override;
    void postprocess() override;
    void decode_outputs(const float *prob, s_detections &objects,
                              int stride, int num_out, const int img_w,
                              const int img_h);
    void priority_filtering(int n, float* crazy, int id);
    s_armor tracker[10];
    s_armor tracker_right[10];
    s_armor zro;
    //std::condition_variable left_cond;
    //std::condition_variable right_cond;
    int left;
    int right;
    std::map<int, std::vector<int> > prior_;
    int now_pri_;
    int simple_nms(float* decode, int n, std::vector<s_armor>& output_nms);

public:
    std::vector<s_armor> output_nms_;
    std::vector<s_armor> output_nms2_;
    Nv_Detector(){};
    ~Nv_Detector(){};
    void Nv_Init();
    bool detect();

};
inline float sigmoid(float x) { return (1.0 / (1.0 + exp(-x))); }

/**
 * @brief calculate sigmoid values in an array
 * 
 * @param src pointer of source array
 * @param dst pointer of destination array
 * @param length number of values
 */
inline void sigmoid(const float *src, float *dst, int length) {
    for (int i = 0; i < length; ++i) {
        dst[i] = (1.0 / (1.0 + exp(-src[i])));
    }
}
#else 
#endif

#endif
